#include <iostream>
#include <fstream>
#include <vector>
#include <cuda_runtime.h>
#include "NvInfer.h"
#include "NvOnnxConfig.h"
#include "NvInferRuntime.h"

#include <algorithm>
#include <math.h>
#include <sys/time.h>
#include <unistd.h>
#include <chrono>

#include "dsgn.hpp"
// #include "common/"
#include "common/check.hpp"
#include "common/timer.hpp"
#include "postprocess.hpp"

namespace dsgn
{
    
class DsgnCoreImpl : public Core {
    public:
    // virtual ~DsgnCoreImpl() {
    // }
    bool init(const CoreParameter& param){

        printf("Initializing DsgnCoreImpl\n");
        normalizer_ = create_normalization(param.normalize);
        if (normalizer_ == nullptr) {
            return false;
        }
        bev_model_ = create_bevmodel(param.bev_model);
        if (bev_model_ == nullptr) {
            printf("Failed to create bev_model.\n");
            return false;
        }

        camera_geometry_ = create_geometry(param.geometry);
        if (camera_geometry_ == nullptr) {
        printf("Failed to create geometry.\n");
        return false;
        }

        post_process_ = create_postprocess(param.postprocess);
        if (post_process_ == nullptr) {
            printf("Failed to create post_process.\n");
            return false;
        }

        return true;
    }

    std::vector<Bndbox> forward_only(const unsigned char** images, void* stream){
        // forward pass
        cudaStream_t _stream = static_cast<cudaStream_t>(stream);

        const float* normed_images = this->normalizer_->forward(images, _stream);

        cudaDeviceSynchronize();  // 异步传输
        auto t1 = std::chrono::high_resolution_clock::now();
        this->bev_model_->forward(
            normed_images, 
            this->camera_geometry_->downsampled_disp(),
            this->camera_geometry_->norm_coord_imgs(),
            this->camera_geometry_->valids(),
            _stream);
        
        printf("Forward pass complete\n");
        cudaDeviceSynchronize();  // 异步传输
        auto t2 = std::chrono::high_resolution_clock::now();
        auto time_used = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1);
        std::cout << "Model forward time: " << time_used.count() / 1000 << " ms" << std::endl;

        return this->post_process_->forward(this->bev_model_->get_cls_preds(),
                                     this->bev_model_->get_box_preds(),
                                     this->bev_model_->get_dir_cls_preds(),
                                     _stream);


    }

    std::vector<Bndbox> forward_timer(const unsigned char** images, void* stream){
        //TODO
        return std::vector<Bndbox>();
    }

    virtual std::vector<Bndbox> forward(const unsigned char** images, void* stream) override {
    if (enable_timer_) {
      return this->forward_timer(images, stream);
    } else {
      return this->forward_only(images, stream);
    }
    }
    virtual void set_timer(bool enable) override {
      enable_timer_ = enable;
    }
    virtual void print() override {
        // TODO
        // backbone_.->print();
        // volume_fusion_.->print();
        // head_.->print();
    }
    virtual void update(const float* camera2lidar,  const float* camera_intrinsics, const float* lidar2image,
                      const float* img_aug_matrix, void* stream) override {
        // camera_depth_->update(img_aug_matrix, lidar2image, stream);
        camera_geometry_->update(camera2lidar, camera_intrinsics, img_aug_matrix, stream);
    }
    // virtual void free_excess_memory() override { camera_geometry_->free_excess_memory(); }



    private:
        CoreParameter param_;
        nv::EventTimer timer_;
        std::shared_ptr<Normalization<float>> normalizer_;
        std::shared_ptr<Backbone> backbone_;
        std::shared_ptr<BevModel> bev_model_;
        std::shared_ptr<PostProcessCuda> post_process_;
        std::shared_ptr<Geometry> camera_geometry_;
        // std::shared_ptr<PostProcessCuda> post_process_;

        bool enable_timer_ = false;
        float confidence_threshold_ = 0;

};

std::shared_ptr<Core> create_core(const CoreParameter& param) {
    std::shared_ptr<DsgnCoreImpl> instance(new DsgnCoreImpl());
    printf("Creating DsgnCoreImpl\n");
    if (!instance->init(param)) {
        instance.reset();
        printf("Failed to create DsgnCoreImpl\n");
    }
    return instance;
} 

} // namespace dsgn
